#!/bin/bash

# Create Self-Managed EKS Node Group Script
# This script creates a new self-managed node group based on existing configuration

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Default values
SOURCE_CONFIG_DIR=""
CLUSTER_NAME=""
REGION=""
SOURCE_ASG_NAME=""
NEW_ASG_NAME=""
NEW_LAUNCH_TEMPLATE_NAME=""
NEW_NODE_NAME=""
NEW_AMI_ID=""
NEW_INSTANCE_TYPE=""
MIN_SIZE=""
MAX_SIZE=""
DESIRED_SIZE=""
PROFILE=""
DRY_RUN="false"
LIST_ASGS="false"

# Function to show usage
usage() {
    cat << EOF
Usage: $0 -s SOURCE_CONFIG_DIR -c CLUSTER_NAME -r REGION [OPTIONS]

Required:
  -s, --source-config   Source configuration directory (from collect script)
  -c, --cluster         EKS cluster name
  -r, --region          AWS region

Source Selection:
  --source-asg          Source ASG name to clone (if not specified, will prompt to select)

Customization Options:
  -a, --asg-name        New ASG name (default: original-name-v2)
  -l, --launch-template New launch template name (default: original-name-v2)
  -n, --node-name       New node name prefix (default: original-name-v2)
  -i, --ami-id          New AMI ID (default: latest EKS-optimized AMI)
  -t, --instance-type   New instance type (default: keep original)

Scaling Options:
  --min-size            Minimum size (default: keep original)
  --max-size            Maximum size (default: keep original)
  --desired-size        Desired size (default: keep original)

Other Options:
  -p, --profile         AWS profile name
  --list-asgs           List available ASGs in config and exit
  --dry-run             Show what would be created without actually creating
  -h, --help            Show this help message

Examples:
  # List available ASGs in config
  $0 -s ./eks-nodegroup-config-20250816-194323 --list-asgs

  # Interactive selection (will prompt to choose ASG)
  $0 -s ./eks-nodegroup-config-20250816-194323 -c test-eks-v133 -r us-west-2

  # Clone specific ASG
  $0 -s ./eks-nodegroup-config-20250816-194323 -c test-eks-v133 -r us-west-2 \\
     --source-asg my-nodegroup-asg

  # Create with custom instance type and scaling
  $0 -s ./eks-nodegroup-config-20250816-194323 -c test-eks-v133 -r us-west-2 \\
     --source-asg my-nodegroup-asg -t t3.large --min-size 2 --max-size 5 --desired-size 3

  # Dry run to see what would be created
  $0 -s ./eks-nodegroup-config-20250816-194323 -c test-eks-v133 -r us-west-2 --dry-run
EOF
}

# Function to list available ASGs
list_available_asgs() {
    local asg_dir="$SOURCE_CONFIG_DIR/asgs"
    local -a asg_names=()
    
    if [[ ! -d "$asg_dir" ]]; then
        print_error "ASG directory not found: $asg_dir"
        exit 1
    fi
    
    # Find all ASG JSON files (excluding policies and launch template files)
    for asg_file in "$asg_dir"/*.json; do
        if [[ -f "$asg_file" && ! "$asg_file" =~ (policies|launch-template) ]]; then
            local asg_name=$(basename "$asg_file" .json)
            asg_names+=("$asg_name")
        fi
    done
    
    if [[ ${#asg_names[@]} -eq 0 ]]; then
        print_error "No ASG configuration files found in $asg_dir"
        exit 1
    fi
    
    printf '%s\n' "${asg_names[@]}"
}

# Function to select ASG interactively
select_source_asg() {
    local -a asg_names
    mapfile -t asg_names < <(list_available_asgs)
    
    if [[ ${#asg_names[@]} -eq 1 ]]; then
        print_info "Found single ASG: ${asg_names[0]}"
        echo "${asg_names[0]}"
        return 0
    fi
    
    print_info "Found ${#asg_names[@]} ASGs in configuration:"
    echo ""
    for i in "${!asg_names[@]}"; do
        echo "  $((i+1)). ${asg_names[i]}"
    done
    echo ""
    
    while true; do
        read -p "Select ASG to clone (1-${#asg_names[@]}): " selection
        if [[ "$selection" =~ ^[0-9]+$ ]] && [[ "$selection" -ge 1 ]] && [[ "$selection" -le ${#asg_names[@]} ]]; then
            echo "${asg_names[$((selection-1))]}"
            return 0
        else
            print_error "Invalid selection. Please enter a number between 1 and ${#asg_names[@]}"
        fi
    done
}

# Function to find source ASG configuration
find_source_asg() {
    local asg_dir="$SOURCE_CONFIG_DIR/asgs"
    local selected_asg=""
    
    if [[ ! -d "$asg_dir" ]]; then
        print_error "ASG directory not found: $asg_dir"
        exit 1
    fi
    
    # If source ASG is specified, validate it exists
    if [[ -n "$SOURCE_ASG_NAME" ]]; then
        local asg_file="$asg_dir/$SOURCE_ASG_NAME.json"
        if [[ -f "$asg_file" ]]; then
            echo "$asg_file"
            return 0
        else
            print_error "Specified source ASG configuration not found: $asg_file"
            print_info "Available ASGs:"
            list_available_asgs | sed 's/^/  /'
            exit 1
        fi
    fi
    
    # Interactive selection
    selected_asg=$(select_source_asg)
    local asg_file="$asg_dir/$selected_asg.json"
    
    if [[ -f "$asg_file" ]]; then
        echo "$asg_file"
        return 0
    else
        print_error "Selected ASG configuration file not found: $asg_file"
        exit 1
    fi
}

# Function to find source launch template configuration
find_source_launch_template() {
    local asg_dir="$SOURCE_CONFIG_DIR/asgs"
    local asg_name="$1"
    
    local lt_file="$asg_dir/$asg_name-launch-template-version.json"
    if [[ -f "$lt_file" ]]; then
        echo "$lt_file"
        return 0
    fi
    
    print_error "Launch template configuration not found: $lt_file"
    exit 1
}

# Function to get latest EKS-optimized AMI
get_latest_eks_ami() {
    local eks_version="$1"
    
    print_info "Getting latest EKS-optimized AMI for Kubernetes $eks_version..."
    
    local ami_id=$(aws ec2 describe-images $AWS_CLI_OPTS \
        --owners amazon \
        --filters "Name=name,Values=amazon-eks-node-$eks_version-v*" \
                  "Name=state,Values=available" \
        --query 'Images | sort_by(@, &CreationDate) | [-1].ImageId' \
        --output text 2>/dev/null)
    
    if [[ -n "$ami_id" && "$ami_id" != "None" ]]; then
        echo "$ami_id"
    else
        print_warning "Could not find latest EKS AMI, will use original AMI"
        echo ""
    fi
}

# Function to extract configuration from source files
extract_source_config() {
    local source_asg_file="$1"
    local source_lt_file="$2"
    
    # Extract ASG configuration
    ORIGINAL_ASG_NAME=$(jq -r '.AutoScalingGroups[0].AutoScalingGroupName' "$source_asg_file")
    ORIGINAL_MIN_SIZE=$(jq -r '.AutoScalingGroups[0].MinSize' "$source_asg_file")
    ORIGINAL_MAX_SIZE=$(jq -r '.AutoScalingGroups[0].MaxSize' "$source_asg_file")
    ORIGINAL_DESIRED_SIZE=$(jq -r '.AutoScalingGroups[0].DesiredCapacity' "$source_asg_file")
    ORIGINAL_SUBNETS=$(jq -r '.AutoScalingGroups[0].VPCZoneIdentifier' "$source_asg_file")
    
    # Extract Launch Template configuration
    ORIGINAL_LT_NAME=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateName' "$source_lt_file")
    ORIGINAL_AMI_ID=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateData.ImageId' "$source_lt_file")
    ORIGINAL_INSTANCE_TYPE=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateData.InstanceType' "$source_lt_file")
    ORIGINAL_IAM_PROFILE=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateData.IamInstanceProfile.Name' "$source_lt_file")
    ORIGINAL_USER_DATA=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateData.UserData' "$source_lt_file")
    ORIGINAL_SECURITY_GROUPS=$(jq -r '.LaunchTemplateVersions[0].LaunchTemplateData.SecurityGroupIds | join(",")' "$source_lt_file")
    
    # Set defaults if not provided
    [[ -z "$NEW_ASG_NAME" ]] && NEW_ASG_NAME="${ORIGINAL_ASG_NAME}-v2"
    [[ -z "$NEW_LAUNCH_TEMPLATE_NAME" ]] && NEW_LAUNCH_TEMPLATE_NAME="${ORIGINAL_LT_NAME}-v2"
    [[ -z "$NEW_NODE_NAME" ]] && NEW_NODE_NAME="${ORIGINAL_ASG_NAME}-v2-node"
    [[ -z "$NEW_INSTANCE_TYPE" ]] && NEW_INSTANCE_TYPE="$ORIGINAL_INSTANCE_TYPE"
    [[ -z "$MIN_SIZE" ]] && MIN_SIZE="$ORIGINAL_MIN_SIZE"
    [[ -z "$MAX_SIZE" ]] && MAX_SIZE="$ORIGINAL_MAX_SIZE"
    [[ -z "$DESIRED_SIZE" ]] && DESIRED_SIZE="$ORIGINAL_DESIRED_SIZE"
    
    # Get EKS version for AMI lookup
    if [[ -z "$NEW_AMI_ID" ]]; then
        local cluster_details_file="$SOURCE_CONFIG_DIR/cluster-details.json"
        if [[ -f "$cluster_details_file" ]]; then
            local eks_version=$(jq -r '.cluster.version' "$cluster_details_file")
            NEW_AMI_ID=$(get_latest_eks_ami "$eks_version")
            if [[ -z "$NEW_AMI_ID" ]]; then
                NEW_AMI_ID="$ORIGINAL_AMI_ID"
                print_warning "Using original AMI ID: $NEW_AMI_ID"
            else
                print_success "Found latest EKS AMI: $NEW_AMI_ID"
            fi
        else
            NEW_AMI_ID="$ORIGINAL_AMI_ID"
            print_warning "Using original AMI ID: $NEW_AMI_ID"
        fi
    fi
}

# Function to create launch template
create_launch_template() {
    print_info "Creating launch template: $NEW_LAUNCH_TEMPLATE_NAME"
    
    # Create launch template data JSON
    local lt_data=$(cat << EOF
{
  "IamInstanceProfile": {
    "Name": "$ORIGINAL_IAM_PROFILE"
  },
  "ImageId": "$NEW_AMI_ID",
  "InstanceType": "$NEW_INSTANCE_TYPE",
  "SecurityGroupIds": [$(echo "$ORIGINAL_SECURITY_GROUPS" | sed 's/,/","/g' | sed 's/^/"/' | sed 's/$/"/')],
  "UserData": "$ORIGINAL_USER_DATA",
  "TagSpecifications": [
    {
      "ResourceType": "instance",
      "Tags": [
        {
          "Key": "Name",
          "Value": "$NEW_NODE_NAME"
        },
        {
          "Key": "kubernetes.io/cluster/$CLUSTER_NAME",
          "Value": "owned"
        }
      ]
    }
  ],
  "MetadataOptions": {
    "HttpTokens": "required",
    "HttpPutResponseHopLimit": 2
  }
}
EOF
)
    
    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "DRY RUN: Would create launch template with data:"
        echo "$lt_data" | jq .
        return 0
    fi
    
    aws ec2 create-launch-template $AWS_CLI_OPTS \
        --launch-template-name "$NEW_LAUNCH_TEMPLATE_NAME" \
        --launch-template-data "$lt_data" > /dev/null
    
    print_success "✓ Launch template created: $NEW_LAUNCH_TEMPLATE_NAME"
}

# Function to create Auto Scaling Group
create_auto_scaling_group() {
    print_info "Creating Auto Scaling Group: $NEW_ASG_NAME"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "DRY RUN: Would create ASG with:"
        echo "  Name: $NEW_ASG_NAME"
        echo "  Launch Template: $NEW_LAUNCH_TEMPLATE_NAME"
        echo "  Min/Max/Desired: $MIN_SIZE/$MAX_SIZE/$DESIRED_SIZE"
        echo "  Subnets: $ORIGINAL_SUBNETS"
        return 0
    fi
    
    aws autoscaling create-auto-scaling-group $AWS_CLI_OPTS \
        --auto-scaling-group-name "$NEW_ASG_NAME" \
        --launch-template "LaunchTemplateName=$NEW_LAUNCH_TEMPLATE_NAME,Version=\$Latest" \
        --min-size "$MIN_SIZE" \
        --max-size "$MAX_SIZE" \
        --desired-capacity "$DESIRED_SIZE" \
        --vpc-zone-identifier "$ORIGINAL_SUBNETS" \
        --tags "Key=Name,Value=$NEW_ASG_NAME,PropagateAtLaunch=true" \
               "Key=kubernetes.io/cluster/$CLUSTER_NAME,Value=owned,PropagateAtLaunch=true"
    
    print_success "✓ Auto Scaling Group created: $NEW_ASG_NAME"
}

# Function to display summary
display_summary() {
    echo ""
    print_info "=== Configuration Summary ==="
    echo "Source Configuration: $SOURCE_CONFIG_DIR"
    echo "Cluster Name: $CLUSTER_NAME"
    echo "Region: $REGION"
    echo ""
    echo "Original Configuration:"
    echo "  ASG Name: $ORIGINAL_ASG_NAME"
    echo "  Launch Template: $ORIGINAL_LT_NAME"
    echo "  AMI ID: $ORIGINAL_AMI_ID"
    echo "  Instance Type: $ORIGINAL_INSTANCE_TYPE"
    echo "  Min/Max/Desired: $ORIGINAL_MIN_SIZE/$ORIGINAL_MAX_SIZE/$ORIGINAL_DESIRED_SIZE"
    echo ""
    echo "New Configuration:"
    echo "  ASG Name: $NEW_ASG_NAME"
    echo "  Launch Template: $NEW_LAUNCH_TEMPLATE_NAME"
    echo "  Node Name: $NEW_NODE_NAME"
    echo "  AMI ID: $NEW_AMI_ID"
    echo "  Instance Type: $NEW_INSTANCE_TYPE"
    echo "  Min/Max/Desired: $MIN_SIZE/$MAX_SIZE/$DESIRED_SIZE"
    echo ""
    echo "Reused Configuration:"
    echo "  IAM Instance Profile: $ORIGINAL_IAM_PROFILE"
    echo "  Security Groups: $ORIGINAL_SECURITY_GROUPS"
    echo "  Subnets: $ORIGINAL_SUBNETS"
    echo ""
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -s|--source-config)
            SOURCE_CONFIG_DIR="$2"
            shift 2
            ;;
        -c|--cluster)
            CLUSTER_NAME="$2"
            shift 2
            ;;
        -r|--region)
            REGION="$2"
            shift 2
            ;;
        --source-asg)
            SOURCE_ASG_NAME="$2"
            shift 2
            ;;
        -a|--asg-name)
            NEW_ASG_NAME="$2"
            shift 2
            ;;
        -l|--launch-template)
            NEW_LAUNCH_TEMPLATE_NAME="$2"
            shift 2
            ;;
        -n|--node-name)
            NEW_NODE_NAME="$2"
            shift 2
            ;;
        -i|--ami-id)
            NEW_AMI_ID="$2"
            shift 2
            ;;
        -t|--instance-type)
            NEW_INSTANCE_TYPE="$2"
            shift 2
            ;;
        --min-size)
            MIN_SIZE="$2"
            shift 2
            ;;
        --max-size)
            MAX_SIZE="$2"
            shift 2
            ;;
        --desired-size)
            DESIRED_SIZE="$2"
            shift 2
            ;;
        -p|--profile)
            PROFILE="$2"
            shift 2
            ;;
        --list-asgs)
            LIST_ASGS="true"
            shift
            ;;
        --dry-run)
            DRY_RUN="true"
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            print_error "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Validate required parameters for other operations (skip for list-asgs)
if [[ "$LIST_ASGS" != "true" ]]; then
    if [[ -z "$SOURCE_CONFIG_DIR" || -z "$CLUSTER_NAME" || -z "$REGION" ]]; then
        print_error "Source config directory, cluster name, and region are required"
        usage
        exit 1
    fi
fi

# Check if source config directory exists
if [[ ! -d "$SOURCE_CONFIG_DIR" ]]; then
    print_error "Source config directory not found: $SOURCE_CONFIG_DIR"
    exit 1
fi

# Handle list-asgs option early (before other setup)
if [[ "$LIST_ASGS" == "true" ]]; then
    print_info "Available ASGs in $SOURCE_CONFIG_DIR:"
    list_available_asgs | sed 's/^/  /'
    exit 0
fi

# Set AWS CLI profile if provided
AWS_CLI_OPTS="--region $REGION"
if [[ -n "$PROFILE" ]]; then
    AWS_CLI_OPTS="$AWS_CLI_OPTS --profile $PROFILE"
fi

print_info "Creating new self-managed node group based on: $SOURCE_CONFIG_DIR"
print_info "Cluster: $CLUSTER_NAME"
print_info "Region: $REGION"
print_info "Profile: ${PROFILE:-default}"
if [[ "$DRY_RUN" == "true" ]]; then
    print_warning "DRY RUN MODE - No resources will be created"
fi

# Main execution
main() {
    # Test AWS CLI access
    print_info "Testing AWS CLI access..."
    if ! aws sts get-caller-identity $AWS_CLI_OPTS &>/dev/null; then
        print_error "Cannot access AWS CLI. Please check your credentials and permissions."
        exit 1
    fi
    print_success "✓ AWS CLI access confirmed"
    
    # Find source configuration files
    local source_asg_file=$(find_source_asg)
    local asg_basename=$(basename "$source_asg_file" .json)
    local source_lt_file=$(find_source_launch_template "$asg_basename")
    
    print_success "✓ Selected source ASG: $asg_basename"
    print_success "✓ Found source ASG config: $(basename "$source_asg_file")"
    print_success "✓ Found source launch template config: $(basename "$source_lt_file")"
    
    # Extract configuration from source files
    print_info "Extracting configuration from source files..."
    extract_source_config "$source_asg_file" "$source_lt_file"
    
    # Display summary
    display_summary
    
    if [[ "$DRY_RUN" == "true" ]]; then
        print_warning "DRY RUN MODE - No resources were created"
        exit 0
    fi
    
    # Confirm before creating
    echo ""
    read -p "Do you want to create the new node group? (y/N): " -n 1 -r
    echo ""
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_info "Operation cancelled"
        exit 0
    fi
    
    # Create resources
    create_launch_template
    create_auto_scaling_group
    
    print_success "New self-managed node group created successfully!"
    print_info "Monitor the new nodes joining the cluster with: kubectl get nodes -w"
}

# Run main function
main "$@"